include("TXG_Band_Structure_EfieldScreening_simple.jl")
using LsqFit

# th_deg = 1.08
# Nlayer = 2

# th_deg = 1.57
# Nlayer = 3

# th_deg = 1.77
# Nlayer = 4

th_deg = 1.81
Nlayer = 5

# th_deg = 1.99
# Nlayer = 6

th = th_deg * pi / 180
Efield = 1e8
w = [0.08, 0.1]
config = [0,1,0,1,0,1]
shift = [(0,0),(0,0),(0,0),(0,0),(0,0),(0,0)]
anisotropy = 1
res = 10
Bxlist = range(0, 5, length=6)
bin=0.0002
@show bw = 20e-3 # meV
@show density = 8/sqrt(3)*(th/0.246e-9)^2
@show Cq = density/bw*(1.602e-19)/Nlayer

pyplot()

folder = mkpath("txg_$(Nlayer)_$(th_deg)_E$(Efield)_w$(w[1])_$(w[2])$(anisotropy==1 ? "" : "_ani$(anisotropy)")_screened_Cq$(Cq)")

deltaEp = zeros(length(Bxlist))
deltaEn = zeros(length(Bxlist))
avg(x) = sum(x[:])/length(x[:])

@time for i=1:length(Bxlist)
    Bx=Bxlist[i]
    (kx,ky,E,Ep,Dirac) = calcTXGBandsFullBZ_screened(
        Nlayer,
        config,
        th,
        res,
        shift,
        Efield=Efield,
        Cq=Cq,
        title="theta=$(th_deg)deg, E=$(Efield)",
        Bin=[Bx,0],
        w=w,
        anisotropy=anisotropy,
        save="$(folder)/bands_Bx$(Bx).txt"
        )

    deltaEp[i] = avg(abs.(Ep[Dirac+1,:,:]-E[Dirac+1,:,:]))
    deltaEn[i] = avg(abs.(Ep[Dirac,:,:]-E[Dirac,:,:]))
    # calcDOS(
    #     kx, ky, E, res, bin, th,
    #     save="$(folder)/dos_Bx$(Bx).txt"
    # )
end


h=plot(Bxlist, deltaEn);
plot!(Bxlist, deltaEp);
display(h)


fp = open("$(folder)/deltaE.txt", "w")
for i=1:length(Bxlist)
    println(fp, "$(Bxlist[i])\t$(deltaEp[i])\t$(deltaEn[i])")
end
close(fp)

# Get g-factor by linear fit
model(x,p) = p[1].+p[2].*x
fp = curve_fit(model, Bxlist|>collect, deltaEp, [0.0,0.0])
fn = curve_fit(model, Bxlist|>collect, deltaEn, [0.0,0.0])
println("+: g=$(fp.param[2]/0.00005788)")
println("-: g=$(fn.param[2]/0.00005788)")
